package com.ruoyi.event.controller;

import com.ruoyi.ai.domain.ChatModelList;
import com.ruoyi.ai.domain.ChatSession;
import com.ruoyi.ai.domain.ChatSessionDetail;
import com.ruoyi.ai.service.ChatAssistant;
import com.ruoyi.ai.service.IChatModelListService;
import com.ruoyi.ai.service.IChatSessionDetailService;
import com.ruoyi.ai.service.IChatSessionService;
import com.ruoyi.ai.utils.ChatSessionManager;
import com.ruoyi.common.annotation.Anonymous;
import com.ruoyi.common.annotation.Log;
import com.ruoyi.common.core.controller.BaseController;
import com.ruoyi.common.core.domain.AjaxResult;
import com.ruoyi.common.core.page.TableDataInfo;
import com.ruoyi.common.enums.BusinessType;
import com.ruoyi.event.service.WxLoginService;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.util.Date;
import java.util.List;
import java.util.Map;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@RestController
@RequestMapping("/event/ai")
public class AiChatController extends BaseController {
    private static final Logger log = LoggerFactory.getLogger(AiChatController.class);

    @Autowired
    private IChatSessionService chatSessionService;

    @Autowired
    private ChatSessionManager chatSessionManager;
    /**
     * 新增聊天会话
     */
    @Log(title = "聊天会话", businessType = BusinessType.INSERT)
    @PostMapping("/add/session")
    public AjaxResult add(@RequestBody ChatSession chatSession)
    {
        int row = chatSessionService.insertChatSession(chatSession);
        chatSessionManager.createSession(chatSession);
        return toAjax(row);
    }


    @Autowired
    private IChatSessionDetailService chatSessionDetailService;

    /**
     * 查询聊天会话详情列表
     */
    @GetMapping("/detail/list")
    public TableDataInfo list(ChatSessionDetail chatSessionDetail)
    {
        List<ChatSessionDetail> list = chatSessionDetailService.selectChatSessionDetailList(chatSessionDetail);
        return getDataTable(list);
    }

    @Autowired
    private EmbeddingStore embeddingStore;

    @Autowired
    private InMemoryEmbeddingStore inMemoryEmbeddingStore;
    @Autowired
//    private OllamaStreamingChatModel ollamaStreamingChatModel;
    private StreamingChatLanguageModel streamingChatLanguageModel;  // 改为接口类型

    @Autowired
    private IChatModelListService chatModelListService;

    @Autowired
    private OllamaEmbeddingModel ollamaEmbeddingModel;
//    private EmbeddingModel ollamaEmbeddingModel;

    /**
     * 普通模式聊天（非流式）
     */
    @PostMapping("/completeChat")
    @Anonymous
    public AjaxResult completeChat(@RequestBody Map<String, Object> request) {
        try {
            String text = (String) request.get("text");
            Long sessionId = Long.valueOf(request.get("sessionId").toString());

            // 获取会话内存
            ChatMemory memory = chatSessionManager.getSessionMemory(sessionId);
            if (memory == null) {
                return AjaxResult.error("无效的会话ID");
            }

            log.info(memory.toString());

            // 保存用户消息
            ChatSessionDetail userMsg = new ChatSessionDetail();
            userMsg.setSessionId(sessionId);
            userMsg.setRole("user");
            userMsg.setContent(text);
            userMsg.setCreateTime(new Date());
            chatSessionDetailService.insertChatSessionDetail(userMsg);

            // 构建内容检索器
//            EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
            ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                    .embeddingStore(embeddingStore)
                    .embeddingModel(ollamaEmbeddingModel)
                    .maxResults(2)
                    .minScore(0.8)
                    .build();

            // 获取启用的模型配置
            ChatModelList chooseModel = chatModelListService.selectEnableChatModel();
            if (chooseModel == null) {
                return AjaxResult.error("未配置可用模型");
            }

            log.info(chooseModel.toString());


            // 构建非流式聊天模型
//            OllamaChatModel chatModel = OllamaChatModel.builder()
//                    .baseUrl(chooseModel.getBaseUrl())
//                    .modelName(chooseModel.getModelVersion())
//                    .temperature(chooseModel.getTemperature().doubleValue())
//                    .timeout(Duration.ofSeconds(300))
//                    .topK(chooseModel.getTopK())
//                    .topP(chooseModel.getTopP().doubleValue())
//                    .repeatPenalty(chooseModel.getRepeatPenalty().doubleValue())
//                    .numPredict(chooseModel.getNumPredict())
//                    .numCtx(chooseModel.getNumPredict())
//                    .maxRetries(3)
//                    .build();

            // 替换为 OpenAI 兼容模型（DeepSeek）
            ChatLanguageModel chatModel = OpenAiChatModel.builder()
                    .baseUrl(chooseModel.getBaseUrl())
                    .apiKey(chooseModel.getApiKey())
                    .modelName(chooseModel.getModelVersion())
                    .temperature(chooseModel.getTemperature().doubleValue())
                    .timeout(Duration.ofSeconds(300))
                    .maxRetries(3)
                    .logRequests(true)
                    .logResponses(true)
                    .build();

            // 创建聊天助手
            ChatAssistant chatAssistant = AiServices.builder(ChatAssistant.class)
                    .chatLanguageModel(chatModel)
                    .contentRetriever(contentRetriever)
                    .chatMemory(memory)
                    .build();

            // 同步获取完整响应
            String aiResponse = chatAssistant.chat(text);

            // 保存AI响应
            ChatSessionDetail aiMsg = new ChatSessionDetail();
            aiMsg.setSessionId(sessionId);
            aiMsg.setRole("assistant");
            aiMsg.setContent(aiResponse);
            aiMsg.setCreateTime(new Date());
            chatSessionDetailService.insertChatSessionDetail(aiMsg);

            return AjaxResult.success(Map.of("content", aiResponse));
        } catch (Exception e) {
            log.error("普通模式聊天失败", e);
            return AjaxResult.error("聊天处理失败：" + e.getMessage());
        }
    }


    /**
     * chat
     */
    @GetMapping(value = "/streamChat", produces = "text/event-stream")
    @Anonymous
    public Flux<String> streamChat(@RequestParam("text") String text, @RequestParam("sessionId") Long sessionId) {
        ChatMemory memory = chatSessionManager.getSessionMemory(sessionId);
        if (memory == null) {
            throw new IllegalArgumentException("Invalid session ID");
        }
        // 保存用户消息
        ChatSessionDetail userMsg = new ChatSessionDetail ();
        userMsg.setSessionId(sessionId);
        userMsg.setRole("user");
        userMsg.setContent(text);
        userMsg.setCreateTime(new Date());
        chatSessionDetailService.insertChatSessionDetail(userMsg);
        EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(2) // on each interaction we will retrieve the 2 most relevant segments
                .minScore(0.2) // we want to retrieve segments at least somewhat similar to user query
                .build();
        ChatAssistant chatAssistant = AiServices.builder(ChatAssistant.class)
//                .streamingChatLanguageModel(ollamaStreamingChatModel)
                .streamingChatLanguageModel(streamingChatLanguageModel)
                .contentRetriever(contentRetriever)
                .chatMemory(memory)
                .build();
        StringBuilder responseBuffer = new StringBuilder();
        return chatAssistant.ollamaChatStreaming(text)
                .doOnNext(responseBuffer::append)
                .doFinally(e -> {
                    // 保存AI响应消息
                    ChatSessionDetail aiMsg = new ChatSessionDetail();
                    aiMsg.setSessionId(sessionId);
                    aiMsg.setRole("assistant");
                    aiMsg.setContent(responseBuffer.toString());
                    aiMsg.setCreateTime(new Date());
                    chatSessionDetailService.insertChatSessionDetail(aiMsg);
                });
    }
}
